package main

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"runtime"

	_ "github.com/go-sql-driver/mysql"
	_ "github.com/jackc/pgx/v5/stdlib"
	jsoniter "github.com/json-iterator/go"
	"github.com/uptrace/bun"
	"github.com/uptrace/bun/dialect/mysqldialect"
	"github.com/uptrace/bun/dialect/pgdialect"
	"github.com/uptrace/bun/dialect/sqlitedialect"
	"github.com/uptrace/bun/extra/bundebug"
	"github.com/uptrace/bun/extra/bunjson"
	"github.com/uptrace/bun/schema"
	"github.com/xiusin/logger"
)

var _ bunjson.Provider = (*JsonProvider)(nil)

type JsonProvider struct{}

func (JsonProvider) Marshal(v interface{}) ([]byte, error) {
	return jsoniter.Marshal(v)
}

func (JsonProvider) Unmarshal(data []byte, v interface{}) error {
	return jsoniter.Unmarshal(data, v)
}

func (JsonProvider) NewEncoder(w io.Writer) bunjson.Encoder {
	return jsoniter.NewEncoder(w)
}

func (JsonProvider) NewDecoder(r io.Reader) bunjson.Decoder {
	return jsoniter.NewDecoder(r)
}

type Config struct {
	DriverName string
	Dsn        string
}

func (cfg Config) IsMysql() bool {
	return cfg.DriverName == "mysql"
}

func (cfg Config) IsSqlite() bool {
	return cfg.DriverName == "sqlite"
}

func (cfg Config) IsPG() bool {
	return cfg.DriverName == "pgx" || cfg.DriverName == "pgsql"
}

type DbClient struct {
	db  *bun.DB
	cfg Config
	ctx context.Context
}

func NewDbClient() *DbClient {
	return &DbClient{}
}

func (client *DbClient) SetContext(ctx context.Context) {
	client.ctx = ctx
}

// Connect 连接数据到指定数据库
func (client *DbClient) Connect(cfg Config, tests ...bool) error {
	var err error
	var conn *sql.DB
	client.cfg = cfg

	if cfg.DriverName == "pgsql" {
		cfg.DriverName = "pgx"
	}

	if conn, err = sql.Open(cfg.DriverName, cfg.Dsn); err != nil {
		logger.Print("连接数据库失败", err)
		return err
	}
	conn_, err := conn.Conn(context.Background())
	if err != nil {
		return err
	}
	defer conn_.Close()

	logger.Print("成功!")

	if len(tests) > 0 && tests[0] {
		return nil
	}
	conns := 4 * runtime.GOMAXPROCS(0)

	var dialect_ schema.Dialect

	dialect_ = pgdialect.New()
	if cfg.IsSqlite() {
		dialect_ = sqlitedialect.New()
	} else if cfg.IsMysql() {
		dialect_ = mysqldialect.New()
	}

	client.db = bun.NewDB(conn, dialect_, bun.WithDiscardUnknownColumns())
	bunjson.SetProvider(JsonProvider{})
	client.db.SetMaxOpenConns(conns)
	client.db.SetMaxIdleConns(conns)
	client.db.AddQueryHook(bundebug.NewQueryHook(bundebug.WithVerbose(true), bundebug.FromEnv("DB_DEBUG_LEVEL")))

	return nil
}

func (client *DbClient) Disconnect() {
	_ = client.db.Close()
	client.db = nil
	runtime.GC()
}

func (client *DbClient) row2arrMap(rows *sql.Rows) ([]map[string]interface{}, error) {
	if rows == nil {
		return nil, errors.New("no record")
	}
	columns, _ := rows.Columns()
	columnLength := len(columns)
	cache := make([]interface{}, columnLength)
	for index := range cache {
		var a interface{}
		cache[index] = &a
	}
	var list []map[string]interface{}
	for rows.Next() {
		_ = rows.Scan(cache...)
		item := make(map[string]interface{})
		for i, data := range cache {
			item[columns[i]] = *data.(*interface{})
		}
		list = append(list, item)
	}
	err := rows.Close()
	return list, err
}

func (client *DbClient) Delete(database, table string, where string, binding []any, limit int64) bool {
	_, err := client.db.NewDelete().Table(client.GetTable(database, table)).
		Where(where, binding...).Exec(context.Background())
	return err == nil
}

func (client *DbClient) GetTable(database, table string) string {
	if client.cfg.IsMysql() {
		return fmt.Sprintf("%s.%s", database, table)
	}

	return table
}

func (client *DbClient) Total(database, table string) int64 {
	sql_ := fmt.Sprintf("SELECT COUNT(*) as total FROM %s", client.GetTable(database, table))
	var total = struct {
		Total int64
	}{}
	client.db.NewRaw(sql_).Scan(context.Background(), &total)
	return total.Total
}

// 数据记录请求
func (client *DbClient) Records(database, table string, page int) []map[string]any {
	sql_ := fmt.Sprintf("SELECT * FROM %s LIMIT %d", client.GetTable(database, table), page)
	rows, _ := client.db.DB.Query(sql_)
	records, _ := client.row2arrMap(rows)

	for i, it := range records {
		for k, v := range it {
			switch v := v.(type) {
			case []byte:
				it[k] = string(v)
			case sql.NullInt64:
				if v.Valid {
					it[k] = v.Int64
				} else {
					it[k] = nil
				}
			case sql.NullFloat64:
				if v.Valid {
					it[k] = v.Float64
				} else {
					it[k] = nil
				}
			case sql.NullString:
				if v.Valid {
					it[k] = v.String
				} else {
					it[k] = nil
				}
			case sql.NullBool:
				if v.Valid {
					it[k] = v.Bool
				} else {
					it[k] = nil
				}
			case sql.NullTime:
				if v.Valid {
					it[k] = v.Time
				} else {
					it[k] = nil
				}
			default:
				it[k] = v
			}
		}
		records[i] = it
	}

	return records
}

// Databases 数据库列表
func (client *DbClient) Databases() []string {
	sql := map[string]string{
		"mysql":  "SHOW DATABASES",
		"sqlite": "SELECT name FROM sqlite_master WHERE type='table'",
		"pgx":    "SELECT datname FROM pg_catalog.pg_database",
	}

	var databases = []string{}
	client.db.NewRaw(sql[client.cfg.DriverName]).Scan(context.Background(), &databases)
	return databases
}

func (client *DbClient) Tables(database string) []string {

	sql := map[string]string{
		"mysql":  "SHOW TABLES",
		"sqlite": "SELECT name FROM sqlite_master WHERE type='table'",
		"pgx":    "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'",
	}

	var tables []string

	tx, _ := client.db.Begin()
	if client.cfg.IsMysql() {
		tx.NewRaw(fmt.Sprintf("USE %s", database)).Exec(context.Background())
	}
	tx.NewRaw(sql[client.cfg.DriverName]).Scan(context.Background(), &tables)
	tx.Commit()

	return tables
}

func (client *DbClient) Views(database string) []string {
	sql := map[string]string{
		"mysql":  "SELECT TABLE_NAME as table_name FROM information_schema.VIEWS WHERE TABLE_SCHEMA = '" + database + "'",
		"sqlite": "SELECT name as table_name FROM sqlite_master WHERE type='view'",
		"pgx":    "SELECT table_name FROM information_schema.views WHERE table_schema = 'public'",
	}

	var tables []string

	tx, _ := client.db.Begin()
	if client.cfg.IsMysql() {
		_, _ = tx.NewRaw(fmt.Sprintf("USE %s", database)).Exec(context.Background())
	}
	_ = tx.NewRaw(sql[client.cfg.DriverName]).Scan(context.Background(), &tables)
	_ = tx.Rollback()

	return tables
}

type Field struct {
	// Mysql
	Field    string  `bun:"Field" json:"Field,omitempty"`
	Type     string  `bun:"Type" json:"Type,omitempty"`
	Null     string  `bun:"Null" json:"Null,omitempty"`
	Key      string  `bun:"Key" json:"Key,omitempty"`
	Default  *string `bun:"Default" json:"Default,omitempty"`
	Extra    *string `bun:"Extra" json:"Extra,omitempty"`
	Readonly bool    `bun:"-" json:"readonly,omitempty"`

	// -- Sqlite
	Cid                int     `json:"-"`
	Name               string  `json:"-"`
	SqliteFieldType    string  `bun:"type" json:"-"`
	SqliteFieldNotNull bool    `bun:"notnull" json:"-"`
	DfltValue          *string `json:"-"`
	Pk                 bool    `json:"-"`
}

// ToNormalization 归一化
func (field Field) ToNormalization(cfg *Config) Field {
	if cfg.IsSqlite() {
		if field.SqliteFieldType == "" {
			field.SqliteFieldType = "TEXT"
			field.Readonly = true
		}
		field.Type = field.SqliteFieldType
		field.Field = field.Name
		field.Default = field.DfltValue
		if field.Pk {
			field.Key = "PRI"
		}
		if field.SqliteFieldNotNull {
			field.Null = "yes"
		}
	}

	return field
}

// Fields 表字段列表
func (client *DbClient) Fields(database, name string) []Field {
	sql := map[string]string{
		"mysql":  "DESC %s",
		"sqlite": "PRAGMA table_info(%s)",
		"pgx":    "SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '%s'",
	}

	var fields = []Field{}

	tx, _ := client.db.Begin()
	if client.cfg.IsMysql() {
		tx.NewRaw(fmt.Sprintf("USE %s", database)).Exec(context.Background())
	}
	tx.NewRaw(fmt.Sprintf(sql[client.cfg.DriverName], name)).Scan(context.Background(), &fields)
	tx.Rollback()

	for i, field := range fields {
		fields[i] = field.ToNormalization(&client.cfg)
	}

	return fields
}

// DuplicateTable 复制表
func (client *DbClient) DuplicateTable(database, table, newTable string) error {
	if table == newTable {
		return nil
	}

	sql := map[string]string{
		"mysql":  "CREATE TABLE " + newTable + " AS SELECT * FROM " + client.GetTable(database, table),
		"sqlite": "CREATE TABLE " + newTable + " AS SELECT * FROM " + client.GetTable(database, table),
		"pgx":    "CREATE TABLE " + newTable + " AS TABLE " + client.GetTable(database, table),
	}

	_, err := client.db.NewRaw(sql[client.cfg.DriverName]).Exec(context.Background())
	return err
}

// DropTable 删除表
func (client *DbClient) DropTable(database, table string) error {
	_, err := client.db.NewDropTable().Table(client.GetTable(database, table)).Exec(context.Background())
	return err

}

// TruncateTable 截断表
func (client *DbClient) TruncateTable(database, table string) error {
	_, err := client.db.NewTruncateTable().Table(client.GetTable(database, table)).Exec(context.Background())
	return err
}

// RenameTable 重命名表
func (client *DbClient) RenameTable(database, table, newTable string) error {
	if table == newTable {
		return nil
	}

	sql := map[string]string{
		"mysql":  "RENAME TABLE " + client.GetTable(database, table) + " TO " + newTable,
		"sqlite": "CREATE TABLE " + newTable + " AS SELECT * FROM " + client.GetTable(database, table),
		"pgx":    "CREATE TABLE " + newTable + " AS TABLE " + client.GetTable(database, table),
	}

	_, err := client.db.NewRaw(sql[client.cfg.DriverName]).Exec(context.Background())
	return err
}

// CreateTable 创建表
func (client *DbClient) CreateTable() error {
	return nil
}

// CreateTableSql 建表语句
func (client *DbClient) CreateTableSql() (string, error) {
	sql := map[string]string{
		"mysql":  "SHOW CREATE TABLE %s",
		"sqlite": "SELECT sql FROM sqlite_master WHERE type='table' AND name='%s'",
		"pgx":    "待定 %s",
	}
	var ret string
	_, err := client.db.NewRaw(sql[client.cfg.DriverName]).Exec(context.Background(), &ret)
	return ret, err
}

// Close 退出
func (client *DbClient) Close() {
	_ = client.db.Close()
}
